import os
import argparse
import sys
import multiprocessing
gatk='~/biosoft/gatk-4.1.8.1/gatk'

parser = argparse.ArgumentParser(description='输入参数如下:')
parser.add_argument('--process', '-P', type=str, help='进程数量，必要参数', required=True)
parser.add_argument('--gvcf_folder', '-F', help='存储g.vcf的文件夹 必要参数', required=True)
parser.add_argument('--genome', '-G', help='reference genome 必要参数', required=True)
parser.add_argument('--genome_length', '-L', help='染色体的长度文件 eg. chr01:1-29595527 必要参数', required=True)
parser.add_argument('--output', '-O', help='output fold 必要参数', required=True)
parser.add_argument('--prefix', '-N', help='out vcf file name 必要参数', required=True)
args = parser.parse_args()
out_folder=args.output
gvcf_folder=args.gvcf_folder
chr_length_file=args.genome_length
os.system('mkdir '+out_folder)
#### index for g.vcf
gvcf_str=''
pool = multiprocessing.Pool(len(args.process))
for gvcf in os.listdir(gvcf_folder):
    if gvcf.endswith('g.vcf'):
        command=gatk+ ' IndexFeatureFile -F '+gvcf_folder+'/'+gvcf
        gvcf_str+=' -V '+gvcf_folder+'/'+gvcf
        result = pool.apply_async(os.system, (command,))
        #result = pool.apply_async(print, (command,))
pool.close()
pool.join()
pool = multiprocessing.Pool(len(args.process))
loop_chr_length_file=open(chr_length_file,'r')
vcf_str=''
for line in loop_chr_length_file:
    file_name=out_folder+'/'+line.split(':')[0]+'.list'
    out_file=open(file_name,'w')
    out_file.write(line)
    out_file.close()
    command=gatk + ' CombineGVCFs -R '+args.genome+' '+gvcf_str+' -L '+file_name+' -O '+out_folder+'/'+line.split(':')[0]+'.g.vcf && '
    # print(command)
    # sys.exit()
    command2=gatk+' GenotypeGVCFs -R '+ args.genome + ' -V '+ out_folder+'/'+line.split(':')[0]+'.g.vcf -O ' +out_folder+'/'+line.split(':')[0]+'.vcf'
    result = pool.apply_async(os.system, (command+command2,))
    vcf_str+=' -I '+out_folder+'/'+line.split(':')[0]+'.vcf '
    # result = pool.apply_async(print, (command+command2,))
pool.close()
pool.join()
loop_chr_length_file.close

command=gatk+' GatherVcfs '+vcf_str+' -O '+args.prefix
os.system(command)
# print(command)
### 

